import time, threading
import roslib; roslib.load_manifest('sushi_turntable')
import rospy
import tf

from pr2_python import gripper,torso
from pr2_tasks.exceptions import DetectionError
from pr2_tasks import tasks
from sushi_turntable import estimate_model
from geometry_msgs.msg import PoseStamped,Pose,Point,PoseArray
import pr2_python.geometry_tools as geometry_tools
import numpy as np
import copy
from pr2_python.transform_listener import get_transform_listener
from mobile_manipulation_msgs.msg import ObjectTrajectory
from pr2_python.controller_manager_client import ControllerManagerClient
from arm_navigation_msgs.msg import CollisionObject,Shape

""" Topic and service names """
POINT_PUB = "/point_publisher" # for debugging
OBJTRAJ_PUB = "/object_trajectory"
PREDICT_PUB = "/object_trajectory_visualization"
# frame that the cartesian controller expects
CONTROLLER_FRAME = '/torso_lift_link'

class GraspTurntable:
  def __init__(self,arm_name):
    self.tasks = tasks.Tasks()
    self.arm_name = arm_name
    self.arm_planner = self.tasks._arm_tasks._arm_planners[arm_name]
    self.arm_mover = self.tasks._arm_tasks._arm_mover
    self.estimation_model = estimate_model.EstimateRotationModel()
    self.head = self.tasks._head
    self.torso = torso.Torso()
    self.gripper = gripper.Gripper(arm_name)
    self._tf_broadcaster = tf.TransformBroadcaster()
    self._controller_manager = ControllerManagerClient()
    self.model_id = -1

    # internal state
    self._state_lock = threading.Lock()
    if arm_name == 'right_arm':      
      self._desired_pose_pub = rospy.Publisher('r_cart/command_pose', PoseStamped)
    else:
      self._desired_pose_pub = rospy.Publisher('l_cart/command_pose', PoseStamped)
    self._transform_listener = get_transform_listener()
    self.world_interface = self.tasks._wi
    self.planning_scene_interface = self.tasks._psi
    self.reset()

  def reset(self):
    self.last_time = rospy.Time.now()
    self.period = 0.0
    self._cancel_requested = False
    self.collision_object = CollisionObject()
    self.reset_all_ce()


  def cancel_publish(self):
    '''Request to cancel the publishing.    
    '''
    with self._state_lock:
      self._cancel_requested = True

  def get_grasp(self):
    if not self.estimation_model.find_rotation_model():
      raise DetectionError('Could not find rotation model')
    self.estimation_model.publish_results()
    self.estimation_model.publish_tf()
    rotation_model = self.estimation_model.get_rotation_model()
    self.rotation_model = copy.deepcopy(rotation_model)
    if not rotation_model.model_id:
      raise DetectionError('Could not find model')
    model_id = rotation_model.model_id[0]
    self.model_id = model_id
    object_trajectory = self.estimation_model.get_object_trajectory(model_id,rospy.Time.now(),num_intervals=30)

    # Find the best object pose to grasp
    best_index = self._find_best_index(object_trajectory,self.arm_name)
    if best_index < 0:
      raise LookupError('Wrong best index')
    pose_stamped = copy.deepcopy(object_trajectory.poses[best_index])
    # Get the time
    self.last_time = pose_stamped.header.stamp
    self.period = np.abs((2*np.pi)/rotation_model.speed)

    # Find grasps for this object pose
    all_grasps = self.arm_planner.get_grasps(model_id,pose_stamped)
    if not all_grasps:
      raise DetectionError('Could not find grasps')
    grasps_size = 15
    prefer_overhead_grasp_offset = 0.0
    if len(all_grasps) < grasps_size:
      grasps_size = len(all_grasps)
    good_grasps = all_grasps[:grasps_size]

    # Sort the grasps
    sorted_grasps = []
    sorted_grasps.extend(copy.deepcopy(self._sort_grasps(good_grasps,self.arm_name)))

#    # Sort the grasps
#    grasp_list_unsorted = ObjectTrajectory()
#    for gr in good_grasps:
#      print gr
#      # prefer overhead grasps
##      if gr.grasp_pose.position.z < (pose_stamped.pose.position.z + prefer_overhead_grasp_offset):
##        continue
#      grasp_stamped = PoseStamped()
#      grasp_stamped.header = object_trajectory.poses[best_index].header
#      grasp_stamped.pose = gr.grasp_pose
#      grasp_list_unsorted.poses.append(copy.deepcopy(grasp_stamped))

#    grasp_list = ObjectTrajectory()
#    grasp_list.poses.extend(copy.deepcopy(self._sort_poses(grasp_list_unsorted.poses,self.arm_name)))

    grasp_list = ObjectTrajectory()
    for gr in sorted_grasps:
      grasp_stamped = PoseStamped()
      grasp_stamped.header = object_trajectory.poses[best_index].header
      grasp_stamped.pose = gr.grasp_pose
      grasp_list.poses.append(copy.deepcopy(grasp_stamped))
      
    pre_grasp_translation = Point(-0.12,0,0)
    success = False
    bounding_cylinder = self.estimation_model.get_bounding_cylinder()
    self.world_interface.add_object(bounding_cylinder)

    grasp_index = -1
    for i,pose in enumerate(grasp_list.poses):
      pose_stamped_grasp = copy.deepcopy(pose)
      pose_stamped_pre_grasp = copy.deepcopy(pose_stamped_grasp)      
      transformed_point = geometry_tools.transform_point(pre_grasp_translation,pose_stamped_pre_grasp.pose)
      pose_stamped_pre_grasp.pose.position = copy.deepcopy(transformed_point)

      response = self.arm_planner.get_ik(pose_stamped_grasp,collision_aware=False)
      if response.error_code.val == response.error_code.SUCCESS:
        rospy.loginfo('Found valid ik pose for starting grasp')
        response = self.arm_planner.get_ik(pose_stamped_pre_grasp,collision_aware=True)
        if response.error_code.val == response.error_code.SUCCESS:
          rospy.loginfo('Found valid collision-free ik pose for pre-grasp as well')
          self.estimation_model.visualize_pose(pose_stamped_grasp,marker_id=6000,r=0.5,g=0.5,b=0.5,marker_type='ARROW')
          self.estimation_model.visualize_pose(pose_stamped_pre_grasp,marker_id=6001,r=1.0,g=0.1,b=0.8,marker_type='ARROW')
          self.estimation_model.visualize_pose(pose_stamped,marker_id=6002,r=0.0,g=0.0,b=1.0)
          grasp_index = i
          break
    self.world_interface.remove_collision_object(bounding_cylinder.id)
  
    if grasp_index < 0:
      raise DetectionError('Could not find valid grasp')

    self.estimation_model.visualize_pose(object_trajectory.poses[best_index],marker_id=3000,r=1.0,g=0.0,b=0.0,marker_type='ARROW')
    self.estimation_model.visualize_pose(pose_stamped,marker_id=3001,r=0.0,g=0.0,b=1.0)

    object_shape = self.estimation_model._tableware_detection.get_mesh_from_database(model_id)
    self.collision_object.shapes.append(object_shape)
    self.collision_object.id = 'rotating_object'
    self.collision_object.padding = 0.0
    self.collision_object.poses.append(copy.deepcopy(pose_stamped.pose))
    self.collision_object.header = copy.deepcopy(pose_stamped.header)
    
    grasp_result = copy.deepcopy(sorted_grasps[grasp_index])
    return pose_stamped_grasp,pose_stamped_pre_grasp,pose_stamped,pre_grasp_translation,grasp_result
    
  def _find_best_index(self,object_trajectory,arm_name='right_arm'):
    best_point = Point()
    if arm_name == 'right_arm':
      best_point.x = 0.70
      best_point.y = -0.188
      best_point.z = 0.0
    else:
      best_point.x = 0.70
      best_point.y = 0.188
      best_point.z = 0.0
    best_index = 0
    best_distance = 10000
    print range(len(object_trajectory.poses))
    for i in range(len(object_trajectory.poses)):
      distance = self._euclidean_distance(object_trajectory.poses[i].pose.position,best_point)
      if distance < best_distance:
        best_index = i
        best_distance = distance
    #if best_distance > 0.5:
     # return -1
    return best_index
   
  def _sort_poses(self, poses, arm_name):
    best_point = self._get_best_point(arm_name)      
    poses_and_dists = [(p, self._2d_planar_distance_between(best_point, p)) for p in poses]
    poses_and_dists.sort(key=lambda pair: pair[1])
    poses_in_dist_order = [pair[0] for pair in poses_and_dists]
    return poses_in_dist_order   

  def _get_best_point(self,arm_name):
    best_point = PoseStamped()
    if arm_name == 'right_arm':
      best_point.pose.position.x = 0.70
      best_point.pose.position.y = -2.0
      best_point.pose.position.z = 0.0
    else:
      best_point.pose.position.x = 0.70
      best_point.pose.position.y = 0.188
      best_point.pose.position.z = 0.0
    best_point.pose.position.x = self.rotation_model.axis_origin.x
    best_point.pose.position.y = self.rotation_model.axis_origin.y
    best_point.pose.position.z = 0.0
    return best_point

  def _sort_grasps(self, grasps, arm_name):
    best_point = self._get_best_point(arm_name)      
    grasps_and_dists = [(p, self._2d_planar_distance_between_poses(best_point.pose, p.grasp_pose)) for p in grasps]
    grasps_and_dists.sort(key=lambda pair: pair[1])
    grasps_in_dist_order = [pair[0] for pair in grasps_and_dists]
    return grasps_in_dist_order   

#distance between poses and quaternions in plane
  def _euclidean_distance(self,p1, p2):
    '''
    Returns the Euclidean distance between two points.
    
    **Args:**
    
    **p1, p2 (geometry_msgs.msg.Point):** points
    
    **Returns:**
    The Euclidean distance between p1 and p2 in three dimensions.
    '''
    dx = p2.x - p1.x
    dy = p2.y - p1.y
    return np.sqrt(dx*dx + dy*dy)

#distance between poses and quaternions in plane
  def _2d_planar_distance_between(self,p1, p2):
    '''
    Returns the 2D planar distance between two pose stamped.
    
    **Args:**
    
    **p1, p2 (geometry_msgs.Pose):** points
    
    **Returns:**
    The 2D planar distance between p1 and p2.
    '''
    dx = p2.pose.position.x - p1.pose.position.x
    dy = p2.pose.position.y - p1.pose.position.y
    return np.sqrt(dx*dx + dy*dy)

#distance between poses and quaternions in plane
  def _2d_planar_distance_between_poses(self,p1, p2):
    '''
    Returns the 2D planar distance between two poses.
    
    **Args:**
    
    **p1, p2 (geometry_msgs.Pose):** points
    
    **Returns:**
    The 2D planar distance between p1 and p2.
    '''
    dx = p2.position.x - p1.position.x
    dy = p2.position.y - p1.position.y
    return np.sqrt(dx*dx + dy*dy)

  def wait_for_grasp(self):
    rospy.loginfo('waiting for grasp')
    predicted_time = self.last_time.to_time()
    while True:
      predicted_time = predicted_time + self.period
      print predicted_time
      if predicted_time > (rospy.Time.now().to_time()+2.0):
        predicted_time_ros = rospy.Time(predicted_time)
        break
    wait_time = predicted_time_ros.to_time()-rospy.Time.now().to_time()
    rospy.loginfo('Waiting for ' + str(wait_time) + ' seconds')
    my_object_pose = copy.deepcopy(self.estimation_model.get_object_pose(self.model_id,predicted_time_ros))
    self.estimation_model.visualize_pose(my_object_pose,marker_id=3333,r=0.25,g=0.25,b=0.25)
#    print 'predicted object pose'
#    print my_object_pose
    rospy.sleep(wait_time)

  def publish_to_controller(self,goal_pose,grasp_translation,grasp_time,grasp,rate=25.0):
    start_controllers = []
    start_controllers.append('r_cart')
    stop_controllers = []
    stop_controllers.append('r_arm_controller')
    self._controller_manager.switch_controllers(start_controllers=start_controllers,stop_controllers=stop_controllers)
    publish_control_thread = threading.Thread(target = self._publish_to_controller,kwargs=dict(rate=rate,goal_pose=goal_pose,grasp_translation=grasp_translation,grasp_time=grasp_time,grasp=grasp))
    publish_control_thread.start()
    return

  def _publish_to_controller(self,goal_pose,grasp_translation,grasp_time,rate,grasp):
    r = rospy.Rate(rate)
    # This is the pre grasp goal pose
    my_goal_pose = copy.deepcopy(goal_pose)
    desired_goal_pose = copy.deepcopy(goal_pose)
    start_time = rospy.Time.now()
    approach_end_time = start_time + rospy.Duration(grasp_time)
    grasp_end_time = start_time + rospy.Duration(2.0*grasp_time)
    retreat_end_time = start_time + rospy.Duration(3.0*grasp_time)
    first_grasp = True
    while True:
      if self._check_cancel_publish_requested():
        break
      if rospy.Time.now() >= retreat_end_time:
        break
      elif rospy.Time.now() >= grasp_end_time:
        # We are retreating
        dT = (rospy.Time.now() - grasp_end_time).to_sec()
        if dT < 0.0:
          dT = 0.0
        multiplier = -1.0 + dT/grasp_time
        delta_translation = copy.deepcopy(grasp_translation)
        delta_translation.x = multiplier * delta_translation.x
        delta_translation.y = multiplier * delta_translation.y
        delta_translation.z = multiplier * delta_translation.z
        transformed_point = geometry_tools.transform_point(delta_translation,my_goal_pose.pose)
        desired_goal_pose.pose.position = copy.deepcopy(transformed_point)
        desired_goal_pose.header.stamp = rospy.Time.now()
      elif rospy.Time.now() >= approach_end_time:
        if first_grasp:
          self.move_gripper_to_posture(grasp,action_type='grasp')
          first_grasp = False
        # We are grasping
        multiplier = -1.0
        delta_translation = copy.deepcopy(grasp_translation)
        delta_translation.x = multiplier * delta_translation.x
        delta_translation.y = multiplier * delta_translation.y
        delta_translation.z = multiplier * delta_translation.z
        transformed_point = geometry_tools.transform_point(delta_translation,my_goal_pose.pose)
        desired_goal_pose.pose.position = copy.deepcopy(transformed_point)
        desired_goal_pose.header.stamp = rospy.Time.now()
      else:
        # We are approaching
        dT = (rospy.Time.now() - start_time).to_sec()
        if dT < 0.0:
          dT = 0.0
        multiplier = -dT/grasp_time
        delta_translation = copy.deepcopy(grasp_translation)
        delta_translation.x = multiplier * delta_translation.x
        delta_translation.y = multiplier * delta_translation.y
        delta_translation.z = multiplier * delta_translation.z
        transformed_point = geometry_tools.transform_point(delta_translation,my_goal_pose.pose)
        desired_goal_pose.pose.position = copy.deepcopy(transformed_point)
        desired_goal_pose.header.stamp = rospy.Time.now()
      self._transform_listener.waitForTransform(
        CONTROLLER_FRAME, desired_goal_pose.header.frame_id, desired_goal_pose.header.stamp, rospy.Duration(4.0))
      desired_pose_stamped = self._transform_listener.transformPose(CONTROLLER_FRAME, desired_goal_pose)
      self._desired_pose_pub.publish(desired_pose_stamped)
      r.sleep()
    return

  def _check_cancel_publish_requested(self):
    with self._state_lock:
      return self._cancel_requested or rospy.is_shutdown()

  def close_gripper(self):
    close_gripper_thread = threading.Thread(target = self._close_gripper)
    close_gripper_thread.start()
    return

  def _close_gripper(self):
    self.gripper.close()

  def move_gripper_to_posture(self,grasp,action_type):
    move_gripper_to_posture_thread = threading.Thread(target = self._move_gripper_to_posture,kwargs=dict(grasp=grasp,action_type=action_type))
    move_gripper_to_posture_thread.start()
    return

  def _move_gripper_to_posture(self,grasp,action_type):
    self.gripper.move_to_posture(grasp=grasp,action_type=action_type)

  def publish_to_tf(self,current_pose,rate=25.0):
    publish_tf_thread = threading.Thread(target = self._publish_to_tf,kwargs=dict(rate=rate,current_pose=current_pose))
    publish_tf_thread.start()
    return

  def _publish_to_tf(self,current_pose,rate=5.0):
    r = rospy.Rate(rate)
    first_time = True
    while True:
      if self._check_cancel_publish_requested():
        break
      self._tf_broadcaster.sendTransform((current_pose.pose.position.x,current_pose.pose.position.y,current_pose.pose.position.z),
                                  (current_pose.pose.orientation.x,current_pose.pose.orientation.y,current_pose.pose.orientation.z,current_pose.pose.orientation.w),rospy.Time.now(),"object_pre_grasp",current_pose.header.frame_id)
      r.sleep()
    return

  def reset_ce(self):
    self.world_interface.reset_collider_node()
    self.planning_scene_interface.reset()
    rospy.sleep(2.0)

  def reset_all_ce(self):
    self.world_interface.reset()
    self.planning_scene_interface.reset()
    rospy.sleep(2.0)

  def pickup_object(self,move_arm_num_tries=5):
    self.reset()
    self.estimation_model.reset()
    self.tasks.move_arms_to_side()
    grasp_pose,pre_grasp_pose,object_pose,grasp_translation,grasp_result = self.get_grasp()

    pre_grasp_pose_in_object_frame = PoseStamped()
    pre_grasp_pose_in_object_frame.header.frame_id = 'object_pose'
    pre_grasp_pose_in_object_frame.pose = copy.deepcopy(geometry_tools.inverse_transform_pose(pre_grasp_pose.pose,object_pose.pose))
    self.publish_to_tf(pre_grasp_pose_in_object_frame)

    object_frame_in_wrist = PoseStamped()
    object_frame_in_wrist.header.frame_id = 'r_wrist_roll_link'
    object_frame_in_wrist.pose = copy.deepcopy(geometry_tools.inverse_transform_pose(object_pose.pose,grasp_pose.pose))
    
    self.gripper.open()

    num_tries = 0
    success = False
    bounding_cylinder = self.estimation_model.get_bounding_cylinder()
    self.world_interface.add_object(bounding_cylinder)
    while num_tries < move_arm_num_tries:
      handle = self.arm_mover.move_to_goal('right_arm',pre_grasp_pose)
      if handle.reached_goal():
        rospy.loginfo('Move arm seems to say we have reached goal')
        success = True
        break
      else:
        self.reset_ce()
        num_tries = num_tries + 1

    if success:
      self.wait_for_grasp()
      pre_grasp_pose_in_object_frame.header.stamp = rospy.Time.now()
      self.publish_to_controller(goal_pose=pre_grasp_pose_in_object_frame,grasp_translation=grasp_translation,grasp_time=2.0,grasp=grasp_result)
      rospy.sleep(10.0)
      self.collision_object.poses[0] = object_frame_in_wrist.pose
      self.collision_object.header.frame_id = 'r_wrist_roll_link'
      self.collision_object.header.stamp = rospy.Time.now()
      self.world_interface.add_object(self.collision_object)
      self.world_interface.attach_object_to_gripper(self.arm_name,self.collision_object.id)
      self.cancel_publish()
      self.world_interface.remove_collision_object(bounding_cylinder.id)    
    else:
      self.world_interface.remove_collision_object(bounding_cylinder.id)    
      return False

    self.estimation_model.cancel_publish()
		# put the line below in comment - we don't want to move the arm to the side 
		# but instead hold the object in front
		#self.tasks.move_arms_to_side()
    return True
    #self.reset()


